#!/usr/bin/env python3
"""
Comprehensive Test Suite for the GSI Diagnostic Checker Package.

This script tests all four command-line tools provided by the package:
1. check_gsi_diag
2. analyze_gsi_stdout
3. analyze_enkf_stdout
4. compare_gsi_enkf

It creates temporary dummy input files, runs each tool, verifies the output,
and cleans up afterward.
"""

import subprocess
import os
import struct
from pathlib import Path
import argparse

# Get the root directory of the package (gsi_diag_checker)
PACKAGE_ROOT = Path(__file__).parent.parent.resolve()

# --- Configuration ---
TEST_DIR = PACKAGE_ROOT / "test_run"
DIAG_DIR = PACKAGE_ROOT / "examples" / "data"
OUTPUT_DIR = TEST_DIR / "output"

# Define dummy file names based on the user's context
DUMMY_DIAG_ANL = "diag_obs_setup_anl.2018081212"
DUMMY_DIAG_GES = "diag_obs_setup_ges.2018081212"

# Paths to real stdout files
GSI_STDOUT_FILE = Path("/home/docker/comgsi/tutorial/comGSIv3.7_EnKFv1.3/run/job/basic/stdout")
ENKF_STDOUT_FILE = Path("/home/docker/comgsi/tutorial/comGSIv3.7_EnKFv1.3/run/job-test/basic_enkf/stdout")

VERBOSE = False

def print_header(title):
    """Prints a formatted header."""
    print("\n" + "="*80)
    print(f" {title}")
    print("="*80)

def run_command(command, cwd="."):
    """Runs a command and checks for errors."""
    print(f"Running: {' '.join(command)}")
    result = subprocess.run(
        command,
        capture_output=True,
        text=True,
        cwd=cwd
    )
    if result.returncode != 0:
        print(f"❌ FAILURE: Command failed with exit code {result.returncode}")
        print("--- STDOUT ---")
        print(result.stdout)
        print("--- STDERR ---")
        print(result.stderr)
        return False
    print("✅ Command executed successfully.")
    if VERBOSE:
        if result.stdout:
            print("--- STDOUT ---")
            print(result.stdout)
        if result.stderr:
            print("--- STDERR ---")
            print(result.stderr)
    return True

def check_file_exists(filepath):
    """Checks if a file exists and prints a status message."""
    if not filepath.exists():
        print(f"❌ FAILURE: Expected output file not found: {filepath}")
        return False
    print(f"✅ Found expected file: {filepath}")
    return True

def setup_test_environment(run_enkf_tests=False):
    """Create dummy files and directories for testing."""
    print_header("Setting up Test Environment")

    # Create directories
    DIAG_DIR.mkdir(parents=True, exist_ok=True)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # 1. Create dummy binary diagnostic files only if they don't exist
    diag_ges_path = DIAG_DIR / DUMMY_DIAG_GES
    diag_anl_path = DIAG_DIR / DUMMY_DIAG_ANL

    if not diag_ges_path.exists() or not diag_anl_path.exists():
        print("Dummy diagnostic files not found. Creating them...")
        # 2 records, 20 floats each, big-endian format
        diag_data = struct.pack('>40f',
            # Record 1 (GES)
            1, 1, 40.0, -100.0, 0, 850.0, 0, 10.0, 0, 0, 0.5, 0, 0, 0, 0, 0, 280.0, -1.0, 0, 0,
            # Record 2 (GES)
            1, 1, 42.0, -98.0, 0, 800.0, 0, 12.0, 0, 0, 0.4, 0, 0, 0, 0, 0, 282.0, 0.5, 0, 0
        )
        (DIAG_DIR / DUMMY_DIAG_GES).write_bytes(diag_data)
        (DIAG_DIR / DUMMY_DIAG_ANL).write_bytes(diag_data) # Use same data for simplicity
        print(f"Created dummy diagnostic files in: {DIAG_DIR}")
    else:
        print(f"Using existing diagnostic files in: {DIAG_DIR}")

    # 2. Check for real GSI and EnKF stdout files
    if not GSI_STDOUT_FILE.exists():
        print(f"❌ ERROR: GSI stdout file not found at: {GSI_STDOUT_FILE}")
        exit(1)
    print(f"Using GSI stdout file: {GSI_STDOUT_FILE}")

    if run_enkf_tests:
        if not ENKF_STDOUT_FILE.exists():
            print(f"❌ ERROR: EnKF stdout file not found at: {ENKF_STDOUT_FILE}")
            exit(1)
        print(f"Using EnKF stdout file: {ENKF_STDOUT_FILE}")

    print("✅ Test environment setup complete.")

def test_check_gsi_diag():
    """Test the check_gsi_diag tool."""
    print_header("1. Testing 'check_gsi_diag'")
    success = True
    
    # Test basic run
    cmd1 = [
        "check_gsi_diag",
        str(DIAG_DIR / DUMMY_DIAG_GES),
        str(DIAG_DIR / DUMMY_DIAG_ANL)
    ]
    if not run_command(cmd1): success = False

    # Test with plotting
    plot_output_file = OUTPUT_DIR / f"plot_{DUMMY_DIAG_GES}.png"
    cmd2 = [
        "check_gsi_diag",
        "--plot",
        "--output-dir", str(OUTPUT_DIR),
        str(DIAG_DIR / DUMMY_DIAG_GES)
    ]
    if not run_command(cmd2): success = False
    if not check_file_exists(plot_output_file): success = False
    
    # Test with plotting for anl file
    plot_output_anl_file = OUTPUT_DIR / f"plot_{DUMMY_DIAG_ANL}.png"
    cmd3 = [
        "check_gsi_diag",
        "--plot",
        "--output-dir", str(OUTPUT_DIR),
        str(DIAG_DIR / DUMMY_DIAG_ANL)
    ]
    if not run_command(cmd3): success = False
    if not check_file_exists(plot_output_anl_file): success = False

    return success

def test_analyze_gsi_stdout():
    """Test the analyze_gsi_stdout tool."""
    print_header("2. Testing 'analyze_gsi_stdout'")
    success = True
    json_output_file = OUTPUT_DIR / "gsi_analysis.json"
    
    cmd = [
        "analyze_gsi_stdout",
        str(GSI_STDOUT_FILE),
        "--json", str(json_output_file)
    ]
    if not run_command(cmd): success = False
    if not check_file_exists(json_output_file): success = False

    return success

def test_analyze_enkf_stdout():
    """Test the analyze_enkf_stdout tool."""
    print_header("3. Testing 'analyze_enkf_stdout'")
    success = True
    json_output_file = OUTPUT_DIR / "enkf_analysis.json"
    
    cmd = [
        "analyze_enkf_stdout",
        str(ENKF_STDOUT_FILE),
        "--json", str(json_output_file)
    ]
    if not run_command(cmd): success = False
    if not check_file_exists(json_output_file): success = False
    
    return success

def test_compare_gsi_enkf():
    """Test the compare_gsi_enkf tool."""
    print_header("4. Testing 'compare_gsi_enkf'")
    success = True
    json_output_file = OUTPUT_DIR / "comparison.json"
    
    cmd = [
        "compare_gsi_enkf",
        str(GSI_STDOUT_FILE),
        str(ENKF_STDOUT_FILE),
        "--json", str(json_output_file)
    ]
    if not run_command(cmd): success = False
    if not check_file_exists(json_output_file): success = False
    
    return success

def cleanup():
    """Remove all temporary files and directories."""
    print_header("Cleaning Up Test Environment")
    try:
        import shutil
        shutil.rmtree(TEST_DIR)
        print(f"Removed temporary directory: {TEST_DIR}")
    except OSError as e:
        print(f"Error during cleanup: {e}")

def main():
    """Main execution function."""
    global VERBOSE
    parser = argparse.ArgumentParser(
        description="Run integration tests for gsi_diag_checker."
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Print detailed output for each test case."
    )
    parser.add_argument(
        "--enkf",
        action="store_true",
        help="Run tests for analyze_enkf_stdout and compare_gsi_enkf."
    )
    args = parser.parse_args()
    if args.verbose:
        VERBOSE = True

    cleanup_flag = False
    setup_test_environment(args.enkf)
    
    # Run tests
    results = {
        "check_gsi_diag": test_check_gsi_diag(),
        "analyze_gsi_stdout": test_analyze_gsi_stdout(),
    }
    
    if args.enkf:
        results["analyze_enkf_stdout"] = test_analyze_enkf_stdout()
        results["compare_gsi_enkf"] = test_compare_gsi_enkf()

    # Final summary
    print_header("Test Summary")
    failures = 0
    for name, result in results.items():
        status = "✅ PASSED" if result else "❌ FAILED"
        print(f"- {name}: {status}")
        if not result:
            failures += 1
            
    print("-" * 80)
    if failures == 0:
        print("🎉 All tests passed successfully!")
    else:
        print(f"💔 {failures} test(s) failed.")
    if cleanup_flag:
        cleanup()

    # Exit with a status code indicating success or failure
    if failures > 0:
        exit(1)

if __name__ == "__main__":
    main() 